Batch tokenization in the transformers backend's scoring loop#1286
Open
ErenAta16 wants to merge 3 commits into
Open
Batch tokenization in the transformers backend's scoring loop#1286ErenAta16 wants to merge 3 commits into
ErenAta16 wants to merge 3 commits into
Conversation
tok_encode_pair tokenizes a context and its continuations with a separate tokenizer call per string, so a batch of documents ends up issuing one call per context plus one per continuation. On large benchmarks this turns into thousands of small Python-level calls. Adds tok_encode_pair_batch, which does the same encoding (including the trailing-space handling already covered by move_trailing_context_space) but with exactly two tokenizer calls for the whole batch: one for every context, one for every continuation across every document. This lets a fast tokenizer batch and parallelize the work instead of paying per-call overhead for each string. Not wired into any backend yet; the next commit switches the transformers backend's scoring loop over to it.
…tion Replaces the per-document tok_encode_pair loop in _loglikelihood_tokens with tok_encode_pair_batch, so each mini-batch tokenizes its contexts and continuations with two tokenizer calls instead of one pair of calls per document. Verified against test_loglikelihood_eval and test_loglikelihood_padded_tensors_shapes, which exercise this path end to end with a real model.
Guards the actual optimization, not just its correctness: asserts the transformers backend's scoring loop calls _batch_tok_encode twice per mini-batch (once for contexts, once for continuations) regardless of how many documents or choices it contains, instead of once per document. With 3 documents, a per-document loop would make 6 calls.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #732.
tok_encode_pairtokenizes a context and its continuations with a separate tokenizer call per string, so scoring a mini-batch of documents issues one call per context plus one per continuation — on a large benchmark this turns into thousands of small Python-level calls.Adds
tok_encode_pair_batch, which does the same encoding (including the existing trailing-space handling) but with exactly two tokenizer calls per mini-batch: one for every context, one for every continuation across every document. Wires it into the transformers backend's_loglikelihood_tokensloop, which is the hot path the issue describes.Verified that the batched output is identical to the old per-document calls across several
add_special_tokens/move_trailing_context_spacecombinations, confirmed the existing end-to-end loglikelihood tests still pass against a real model, and added a test asserting the call count itself drops (2 calls per mini-batch instead of 2 per document).